// ==========================================================================
// Eses Image Effect Levels
// ==========================================================================
//
// Description:
// The 'Eses Image Effect Levels' node provides a versatile, image editor style
// levels adjustment tool directly within the ComfyUI interface. It allows for
// precise, interactive control over the tonal range of images and masks,
// complete with a live histogram for visual feedback.
//
// Key Features:
//
// - Interactive Controls:
//   - An interactive preview of the level adjustments in the node.
//   - Features sliders for input Black, Mid, and White points, as well as
//     sliders for the final Output Black and White points.
//   - A live histogram visualizes the tonal distribution of the selected
//     channel.
//
// - Multi-Channel Adjustments:
//   - Apply level adjustments to the combined RGB channels for general
//     corrections.
//   - Isolate adjustments to the individual Red, Green, or Blue channels.
//   - Apply a separate level adjustment directly to an input mask.
//
// - State Serialization:
//   - All level adjustments are saved with the workflow and restored on
//     reload.
//   - The node's state persists after refreshing the browser page.
//
// - Live Preview:
//   - The node displays a preview of the connected image with the level
//     adjustments applied in real-time as you move the sliders.
//
// - Quality of Life Features:
//   - "Set Auto Levels" button to automatically calculate optimal
//      input levels.
//   - A simple file-based preset system for saving and loading adjustment
//     presets
//   - "Reset" buttons to revert the current channel or all channels to
//     default values.
//   - Adjustable auto levels sensitivity for fine-tuned automatic adjustments.
//
// Version: 1.3.1
//
// License: See LICENSE.txt
//
// ==========================================================================


import { app } from "/scripts/app.js";
import { api } from "/scripts/api.js";

const HEADER_HEIGHT = 445;
const PADDING = 10;
const HISTOGRAM_AREA_HEIGHT = 40;
const AREA_SPACING = 8;
const SLIDER_BOX_HEIGHT = 15;
const OUTPUT_SLIDER_AREA_HEIGHT = 20;


// --- PRESET HANDLING START ---

// Global cache for presets to avoid
// fetching them for every single node
let allPresets = null;

async function getPresets() {
    if (allPresets === null) {
        try {
            const response = await api.fetchApi("/eses_levels/get_presets", { cache: "no-store" });
            allPresets = await response.json();
            console.log("[EsesImageEffectLevels] Presets loaded:", allPresets);
        }
        catch (e) {
            console.error("[EsesImageEffectLevels] Failed to load presets:", e);
            allPresets = {};
        }
    }

    return allPresets;
}

// Pre-fetch presets when the script is loaded
getPresets();

// --- PRESET HANDLING END ---



app.registerExtension({
    name: "Eses.EsesImageEffectLevels",

    // --- PRESET HANDLING START ---
    async beforeRegisterNodeDef(nodeType, nodeData, app) {

        // Runs once when the node definition is being registered.
        // Dynamically add the preset names to the combo box definition.
        if (nodeData.name === "EsesImageEffectLevels") {
            const presets = await getPresets();
            const presetNames = ["None", ...Object.keys(presets)];

            // Find the 'preset' input in the node's
            // definition and update its options
            const presetInput = nodeData.input.required.preset;
            if (presetInput) {
                presetInput[0] = presetNames;
            }
        }
    },
    // --- PRESET HANDLING END ---


    nodeCreated(node) {

        if (node.comfyClass === "EsesImageEffectLevels") {
            const originalConfigure = node.configure;
            const originalSerialize = node.serialize;
            
            node.size = [256, 575];
            node.imagePreview = null;
            node.originalImage = null;
            node.originalMask = null;
            node.isManuallyResized = false;
            node.dragging = null;
            node.resetButtonState = 'idle';
            
            // Add state for auto color button
            node.autoColorButtonState = 'idle';
            node.onResize = function () { this.isManuallyResized = true; };

            const defaultLevels = () => ({ black_point: 0.0, white_point: 1.0, mid_point: 1.0, output_black: 0, output_white: 255 });
            const resetAllLevels = () => ({ rgb: defaultLevels(), r: defaultLevels(), g: defaultLevels(), b: defaultLevels(), mask: defaultLevels() });

            node.allLevels = resetAllLevels();
            node.histograms = {};

            let blackWidget, midWidget, whiteWidget, outputBlackWidget, outputWhiteWidget;

            function updateWidgetsFromState() {
                const levels = node.allLevels[node.activeChannel];

                if (!levels)
                    return;

                if (blackWidget)
                    blackWidget.value = parseFloat(levels.black_point.toFixed(2));

                if (midWidget)
                    midWidget.value = parseFloat(levels.mid_point.toFixed(2));

                if (whiteWidget)
                    whiteWidget.value = parseFloat(levels.white_point.toFixed(2));

                if (outputBlackWidget)
                    outputBlackWidget.value = Math.round(levels.output_black);

                if (outputWhiteWidget)
                    outputWhiteWidget.value = Math.round(levels.output_white);
            }

            node.getResetButtonName = (isConfirm = false) => {
                const channelWidget = node.widgets.find(w => w.name === "channel");

                if (node.activeChannel === 'rgb')
                    return isConfirm ? "CONFIRM RESET ALL" : "Reset All Levels";

                const channelName = channelWidget.value;

                return isConfirm ? `CONFIRM RESET ${channelName}` : `Reset ${channelName} Level`;
            };

            node.resetButtonCallback = function () {
                const resetButton = node.widgets.find(w => w.name.startsWith("Reset") || w.name.startsWith("CONFIRM"));

                if (node.resetButtonState === 'armed') {
                    clearTimeout(node.resetTimeout);

                    if (node.activeChannel === 'rgb') {
                        node.allLevels = resetAllLevels();
                    }
                    else {
                        node.allLevels[node.activeChannel] = defaultLevels();
                    }

                    updateWidgetsFromState();
                    node.sendLevelUpdateToPython();
                    resetButton.name = node.getResetButtonName();
                    node.resetButtonState = 'idle';
                }
                else {
                    node.resetButtonState = 'armed';
                    resetButton.name = node.getResetButtonName(true);

                    node.resetTimeout = setTimeout(() => {
                        resetButton.name = node.getResetButtonName();
                        node.resetButtonState = 'idle';
                        app.graph.setDirtyCanvas(true);
                    }, 3000);
                }
                app.graph.setDirtyCanvas(true);
            };

            function setupCustomWidgets() {
                const levelsWidget = node.widgets.find(w => w.name === "all_levels_json");
                const channelWidget = node.widgets.find(w => w.name === "channel");
                
                // Hide the boolean widgets that are only used for triggering
                const autoLevelsWidget = node.widgets.find(w => w.name === "auto_levels");
                const autoColorWidget = node.widgets.find(w => w.name === "auto_color");
                
                if (!levelsWidget || !channelWidget)
                    return;

                levelsWidget.type = "hidden";
                if (autoLevelsWidget) autoLevelsWidget.type = "hidden";
                if (autoColorWidget) autoColorWidget.type = "hidden";
                const channelMap = { "RGB": "rgb", "Red": "r", "Green": "g", "Blue": "b", "Mask": "mask" };
                node.activeChannel = channelMap[channelWidget.value] || "rgb";

                if (!node.widgets.find(w => w.name === "Black")) {
                    blackWidget = node.addWidget("number", "Black", 0.0, (val) => {

                        if (node.allLevels[node.activeChannel]) {
                            node.allLevels[node.activeChannel].black_point = val;
                            node.sendLevelUpdateToPython();
                        }
                    }, { min: 0.0, max: 1.0, step: 0.01, precision: 2 });

                    midWidget = node.addWidget("number", "Mid", 1.0, (val) => {

                        if (node.allLevels[node.activeChannel]) {
                            node.allLevels[node.activeChannel].mid_point = val;
                            node.sendLevelUpdateToPython();
                        }
                    }, { min: 0.01, max: 9.99, step: 0.01, precision: 2 });

                    whiteWidget = node.addWidget("number", "White", 1.0, (val) => {

                        if (node.allLevels[node.activeChannel]) {
                            node.allLevels[node.activeChannel].white_point = val;
                            node.sendLevelUpdateToPython();
                        }
                    }, { min: 0.0, max: 1.0, step: 0.01, precision: 2 });

                    outputBlackWidget = node.addWidget("number", "Output Black", 0, (val) => {

                        if (node.allLevels[node.activeChannel]) {
                            node.allLevels[node.activeChannel].output_black = val;
                            node.sendLevelUpdateToPython();
                        }
                    }, { min: 0, max: 255, step: 1, precision: 0 });

                    outputWhiteWidget = node.addWidget("number", "Output White", 255, (val) => {

                        if (node.allLevels[node.activeChannel]) {
                            node.allLevels[node.activeChannel].output_white = val;
                            node.sendLevelUpdateToPython();
                        }
                    }, { min: 0, max: 255, step: 1, precision: 0 });
                }
                else {
                    blackWidget = node.widgets.find(w => w.name === "Black");
                    midWidget = node.widgets.find(w => w.name === "Mid");
                    whiteWidget = node.widgets.find(w => w.name === "White");
                    outputBlackWidget = node.widgets.find(w => w.name === "Output Black");
                    outputWhiteWidget = node.widgets.find(w => w.name === "Output White");
                }

                // --- PRESET HANDLING START ---

                // Remove old "Copy Settings" button if it exists
                const copySettingsButton = node.widgets.find(w => w.name === "Copy Settings");

                if(copySettingsButton)
                    node.widgets.splice(node.widgets.indexOf(copySettingsButton), 1);

                // Add "Preset Name" text field and "Save Preset" button
                if (!node.widgets.find(w => w.name === "Preset Name")) {
                    const presetNameWidget = node.addWidget("string", "Preset Name", "", {});

                    node.addWidget("button", "Save Preset", null, async () => {
                        const name = presetNameWidget.value;

                        if (!name || !name.trim()) {
                            alert("Please enter a preset name.");
                            return;
                        }

                        try {
                            const resp = await api.fetchApi("/eses_levels/save_preset", {
                                method: "POST",
                                headers: { "Content-Type": "application/json" },
                                body: JSON.stringify({
                                    preset_name: name,
                                    preset_data: node.allLevels,
                                }),
                            });

                            const res = await resp.json();

                            if (res.success) {
                                alert(`Preset '${name}' saved successfully!`);

                                // Refresh the presets in the dropdown
                                allPresets = res.presets;
                                const presetWidget = node.widgets.find(w => w.name === "preset");

                                if(presetWidget) {
                                    presetWidget.options.values = ["None", ...Object.keys(allPresets)];
                                    presetWidget.value = name;
                                }
                                presetNameWidget.value = "";
                            }
                            else {
                                alert(`Error saving preset: ${res.message}`);
                            }
                        }
                        catch (e) {
                            console.error("[EsesImageEffectLevels] Error saving preset:", e);
                            alert(`An error occurred: ${e}`);
                        }
                    });
                }

                const presetWidget = node.widgets.find(w => w.name === "preset");
                if (presetWidget) {
                    presetWidget.callback = async (value) => {

                        if (value === "None")
                            return;

                        const presets = await getPresets();
                        const selectedPreset = presets[value];

                        if (selectedPreset) {
                            
                            // Create a fresh default state
                            const newLevels = resetAllLevels();
                            const loadedLevels = JSON.parse(JSON.stringify(selectedPreset));

                            // Deep merge the loaded preset into the default state.
                            // This ensures compatibility with older presets.
                            for (const channelKey in newLevels) {

                                if (loadedLevels[channelKey]) {
                                    Object.assign(newLevels[channelKey], loadedLevels[channelKey]);
                                }
                            }

                            node.allLevels = newLevels;
                            updateWidgetsFromState();
                            node.sendLevelUpdateToPython();
                            node.setDirtyCanvas(true, true);
                        }
                    };
                }

                // --- PRESET HANDLING END ---

                if (!node.widgets.find(w => w.name === "Set Auto Levels")) {
                    node.addWidget("button", "Set Auto Levels", null, () => {
                        const autoLevelsWidget = node.widgets.find(w => w.name === "auto_levels");

                        if (autoLevelsWidget) {
                            const originalValue = autoLevelsWidget.value;
                            autoLevelsWidget.value = true;
                            app.queuePrompt();
                            setTimeout(() => { autoLevelsWidget.value = originalValue; }, 100);
                        }
                    });
                }

                // Add "Set Auto Color" button
                if (!node.widgets.find(w => w.name === "Set Auto Color")) {
                    
                    node.addWidget("button", "Set Auto Color", null, () => {
                        
                        const autoColorWidget = node.widgets.find(w => w.name === "auto_color");
                        
                        if (autoColorWidget) {
                            const originalValue = autoColorWidget.value;
                            autoColorWidget.value = true;
                            app.queuePrompt();
                        
                            // Reset the value back to false after a short delay
                            setTimeout(() => { autoColorWidget.value = originalValue; }, 100);
                        }
                    });
                }


                const oldResetBtn = node.widgets.find(w => w && w.name && (w.name.startsWith("Reset") || w.name.startsWith("CONFIRM")));

                if (oldResetBtn)
                    node.widgets.splice(node.widgets.indexOf(oldResetBtn), 1);

                const resetButton = node.addWidget("button", node.getResetButtonName(), null, node.resetButtonCallback.bind(node));

                channelWidget.callback = (value) => {
                    node.activeChannel = channelMap[value] || "rgb";
                    resetButton.name = node.getResetButtonName();

                    updateWidgetsFromState();
                    node.updateHistogram();
                    node.sendLevelUpdateToPython();
                };
            }

            node.configure = function (data) {
                originalConfigure.apply(this, arguments);

                if (data.custom_all_levels) {
                    this.allLevels = data.custom_all_levels;
                }

                setupCustomWidgets();
                updateWidgetsFromState();
            };

            node.serialize = function () {
                const data = originalSerialize.call(this);
                data.custom_all_levels = node.allLevels;
                return data;
            };

            setupCustomWidgets();

            Object.assign(node, {

                getImageArea: function () {
                    const interactiveHeight = HISTOGRAM_AREA_HEIGHT + OUTPUT_SLIDER_AREA_HEIGHT + AREA_SPACING;
                    
                    // Moved image top edge down by the 
                    // amount increased by HEADER_HEIGHT
                    const totalContentHeight = this.size[1] - HEADER_HEIGHT - PADDING * 2;
                    const imageAreaHeight = totalContentHeight - interactiveHeight;

                    if (imageAreaHeight < 10)
                        return null;

                    return { x: PADDING, y: HEADER_HEIGHT + PADDING, width: this.size[0] - PADDING * 2, height: imageAreaHeight };
                },

                getBarArea: function () {
                    const imageArea = this.getImageArea();

                    if (!imageArea)
                        return null;

                    const barAreaHeight = HISTOGRAM_AREA_HEIGHT - SLIDER_BOX_HEIGHT - 4;
                    return { x: PADDING, y: imageArea.y + imageArea.height + AREA_SPACING, width: this.size[0] - PADDING * 2, height: barAreaHeight };
                },

                getInputSliderArea: function () {
                    const barArea = this.getBarArea();

                    if (!barArea)
                        return null;

                    return { x: PADDING, y: barArea.y + barArea.height + 4, width: this.size[0] - PADDING * 2, height: SLIDER_BOX_HEIGHT };
                },

                getOutputSliderArea: function () {
                    const inputSliderArea = this.getInputSliderArea();

                    if (!inputSliderArea)
                        return null;

                    return { x: PADDING, y: inputSliderArea.y + inputSliderArea.height + AREA_SPACING, width: this.size[0] - PADDING * 2, height: SLIDER_BOX_HEIGHT };
                },


                updateHistogram: async function () {
                    const channelWidget = node.widgets.find(w => w.name === "channel");
                    let base_data_image = (this.activeChannel === 'mask') ? this.originalMask : this.originalImage;

                    if (!base_data_image)
                        return;

                    this.histograms[this.activeChannel] = null;
                    const canvas = document.createElement('canvas');
                    canvas.width = base_data_image.naturalWidth;
                    canvas.height = base_data_image.naturalHeight;
                    canvas.getContext('2d').drawImage(base_data_image, 0, 0);
                    const base_data_b64 = canvas.toDataURL('image/png').split(',')[1];

                    try {
                        const result = await api.fetchApi("/eses_levels/get_histogram", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ base_data_b64, channel_mode: channelWidget.value }) }).then(r => r.json());

                        if (result.histogram) {
                            this.histograms[this.activeChannel] = result.histogram;
                            this.setDirtyCanvas(true, true);
                        }
                    }
                    catch (err) {
                        console.error("Failed to fetch histogram:", err);
                    }
                },

                sendLevelUpdateToPython: async function () {
                    this.setDirtyCanvas(true, true);

                    const levelsWidget = this.widgets.find(w => w.name === "all_levels_json");

                    if (!levelsWidget)
                        return;

                    levelsWidget.value = JSON.stringify(this.allLevels);
                    let base_data_type, base_data_image;

                    if (this.activeChannel === 'mask') {
                        if (!this.originalMask) {
                            this.imagePreview = null; this.setDirtyCanvas(true, true); return;
                        }

                        base_data_type = 'mask';
                        base_data_image = this.originalMask;
                    }
                    else {
                        if (!this.originalImage) {
                            this.imagePreview = null; this.setDirtyCanvas(true, true); return;
                        }

                        base_data_type = 'image';
                        base_data_image = this.originalImage;
                    }

                    const canvas = document.createElement('canvas');
                    canvas.width = base_data_image.naturalWidth;
                    canvas.height = base_data_image.naturalHeight;
                    canvas.getContext('2d').drawImage(base_data_image, 0, 0);
                    const base_data_b64 = canvas.toDataURL('image/png').split(',')[1];

                    const body = {
                        base_data_b64,
                        all_levels_json: levelsWidget.value,
                        base_data_type
                    };

                    try {
                        const result = await api.fetchApi("/eses_levels/apply", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(body) }).then(r => r.json());

                        if (result.adjusted_image_data) {
                            const img = new Image();
                            img.src = `data:image/png;base64,${result.adjusted_image_data}`;

                            img.onload = () => {
                                this.imagePreview = img; this.setDirtyCanvas(true, true);
                            };
                        }
                    }
                    catch (err) { }
                },

                drawHistogram: function (ctx, area, data, step = 1) {
                    if (!data || data.length === 0)
                        return;

                    step = Math.max(1, Math.floor(step));
                    ctx.fillStyle = "#999";
                    const baseBarWidth = area.width / 256;

                    for (let i = 0; i < 256; i += step) {
                        const chunk = data.slice(i, i + step);

                        if (chunk.length === 0)
                            continue;

                        const maxVal = Math.max(...chunk);
                        const barHeight = maxVal * area.height;
                        const barWidth = baseBarWidth * chunk.length;
                        const x = area.x + i * baseBarWidth;
                        ctx.fillRect(x, area.y + area.height - barHeight, barWidth, barHeight);
                    }
                },

                drawHandle: function (ctx, x, y, size, color, strokeColor) {
                    ctx.fillStyle = color;
                    ctx.strokeStyle = strokeColor;
                    ctx.lineWidth = 1;
                    ctx.beginPath();
                    ctx.moveTo(x, y);
                    ctx.lineTo(x - size, y + size * 2);
                    ctx.lineTo(x + size, y + size * 2);
                    ctx.closePath();
                    ctx.fill();
                    ctx.stroke();
                },


                onDrawForeground: function (ctx) {

                    if (this.flags?.collapsed)
                        return;

                    const imageArea = this.getImageArea();
                    const barArea = this.getBarArea();
                    const inputSliderArea = this.getInputSliderArea();
                    const outputSliderArea = this.getOutputSliderArea();

                    if (!imageArea || !barArea || !inputSliderArea || !outputSliderArea)
                        return;

                    if (this.imagePreview) {
                        ctx.drawImage(this.imagePreview, imageArea.x, imageArea.y, imageArea.width, imageArea.height);
                    }
                    else {
                        ctx.save();
                        ctx.fillStyle = LiteGraph.NODE_TEXT_COLOR || "#CCC";
                        ctx.textAlign = "center";
                        ctx.textBaseline = "middle";
                        ctx.font = "bold 14px Arial";
                        const message = "Connect inputs and run workflow";
                        ctx.fillText(message, imageArea.x + imageArea.width / 2, imageArea.y + imageArea.height / 2);
                        ctx.restore();
                    }

                    // Draw Input Levels Area
                    ctx.fillStyle = "#282828";
                    ctx.fillRect(barArea.x, barArea.y, barArea.width, barArea.height);
                    ctx.fillStyle = "#202020";
                    ctx.fillRect(inputSliderArea.x, inputSliderArea.y, inputSliderArea.width, inputSliderArea.height);

                    const histogramData = this.histograms[this.activeChannel];

                    if (histogramData) {
                        this.drawHistogram(ctx, barArea, histogramData, 1);
                    }

                    const levels = this.allLevels[this.activeChannel];
                    if (!levels) return;

                    // Draw Input Handles
                    const blackX = inputSliderArea.x + levels.black_point * inputSliderArea.width;
                    const whiteX = inputSliderArea.x + levels.white_point * inputSliderArea.width;
                    const relativeRangeWidth = whiteX - blackX;

                    const midPointToSlider = (gamma) => {
                        if (gamma >= 1.0)
                            return 1.0 - (0.5 + ((gamma - 1.0) / 8.99) * 0.5);
                        else return 1.0 - (((gamma - 0.01) / 0.99) * 0.5);
                    };

                    const midHandleX = blackX + (midPointToSlider(levels.mid_point) * relativeRangeWidth);
                    const handleTipY = inputSliderArea.y + 2.5;

                    this.drawHandle(ctx, blackX, handleTipY, 5, "black", "white");
                    this.drawHandle(ctx, whiteX, handleTipY, 5, "white", "black");
                    this.drawHandle(ctx, midHandleX, handleTipY, 5, "#888", "white");

                    // Draw Output Levels Area
                    const gradient = ctx.createLinearGradient(outputSliderArea.x, 0, outputSliderArea.x + outputSliderArea.width, 0);
                    const blackOutPos = levels.output_black / 255.0;
                    const whiteOutPos = levels.output_white / 255.0;

                    gradient.addColorStop(0, "black");
                    gradient.addColorStop(blackOutPos, "black");
                    gradient.addColorStop(whiteOutPos, "white");
                    gradient.addColorStop(1, "white");
                    ctx.fillStyle = gradient;
                    ctx.fillRect(outputSliderArea.x, outputSliderArea.y, outputSliderArea.width, outputSliderArea.height);

                    // Draw Output Handles
                    const outBlackX = outputSliderArea.x + blackOutPos * outputSliderArea.width;
                    const outWhiteX = outputSliderArea.x + whiteOutPos * outputSliderArea.width;
                    const outHandleTipY = outputSliderArea.y + 2.5;

                    this.drawHandle(ctx, outBlackX, outHandleTipY, 5, "black", "white");
                    this.drawHandle(ctx, outWhiteX, outHandleTipY, 5, "white", "black");
                },


                onMouseDown: function (event) {

                    if (event.button !== 0)
                        return false;

                    const localPos = app.canvas.convertEventToCanvasOffset(event);
                    const mouseX = localPos[0] - this.pos[0];
                    const mouseY = localPos[1] - this.pos[1];

                    const levels = this.allLevels[this.activeChannel];

                    if (!levels)
                        return false;

                    // Check output slider first
                    const outputSliderArea = this.getOutputSliderArea();
                    if (outputSliderArea && mouseX > outputSliderArea.x && mouseX < outputSliderArea.x + outputSliderArea.width &&
                        mouseY > outputSliderArea.y && mouseY < outputSliderArea.y + outputSliderArea.height) {

                        const outBlackX = outputSliderArea.x + (levels.output_black / 255.0) * outputSliderArea.width;

                        if (Math.abs(mouseX - outBlackX) < 10) {
                            this.dragging = 'output_black'; return true;
                        }

                        const outWhiteX = outputSliderArea.x + (levels.output_white / 255.0) * outputSliderArea.width;

                        if (Math.abs(mouseX - outWhiteX) < 10) {
                            this.dragging = 'output_white'; return true;
                        }
                    }


                    // Check input slider
                    const inputSliderArea = this.getInputSliderArea();
                    if (inputSliderArea && mouseX > inputSliderArea.x && mouseX < inputSliderArea.x + inputSliderArea.width &&
                        mouseY > inputSliderArea.y && mouseY < inputSliderArea.y + inputSliderArea.height) {

                        const blackX = inputSliderArea.x + levels.black_point * inputSliderArea.width;
                        const whiteX = inputSliderArea.x + levels.white_point * inputSliderArea.width;
                        const relativeRangeWidth = whiteX - blackX;

                        const midPointToSlider = (gamma) => {
                            if (gamma >= 1.0) {
                                return 1.0 - (0.5 + ((gamma - 1.0) / 8.99) * 0.5);
                            }
                            else {
                                return 1.0 - (((gamma - 0.01) / 0.99) * 0.5);
                            }
                        };

                        const midHandleX = blackX + (midPointToSlider(levels.mid_point) * relativeRangeWidth);

                        if (Math.abs(mouseX - midHandleX) < 10) {
                            this.dragging = 'mid';
                        }
                        else if (Math.abs(mouseX - blackX) < 10) {
                            this.dragging = 'black';
                        }
                        else if (Math.abs(mouseX - whiteX) < 10) {
                            this.dragging = 'white';
                        }
                    }

                    return this.dragging !== null;
                },

                onMouseMove: function (event) {
                    if (event.buttons !== 1) {

                        if (this.dragging) {
                            this.dragging = null; this.sendLevelUpdateToPython();
                        }
                        return;
                    }

                    if (!this.dragging)
                        return;

                    const localPos = app.canvas.convertEventToCanvasOffset(event);
                    const mouseX = localPos[0] - this.pos[0];
                    const levels = this.allLevels[this.activeChannel];

                    if (!levels)
                        return;

                    let needsUpdate = false;

                    // Handle input slider dragging
                    if (this.dragging === 'black' || this.dragging === 'white' || this.dragging === 'mid') {
                        const inputSliderArea = this.getInputSliderArea();

                        if (!inputSliderArea)
                            return;

                        if (this.dragging === 'black' || this.dragging === 'white') {
                            let newLevel = (mouseX - inputSliderArea.x) / inputSliderArea.width;
                            newLevel = Math.max(0.0, Math.min(1.0, newLevel));

                            if (this.dragging === 'black')
                                levels.black_point = Math.min(newLevel, levels.white_point - 0.01);
                            else
                                levels.white_point = Math.max(newLevel, levels.black_point + 0.01);
                        }
                        else if (this.dragging === 'mid') {
                            const blackX = inputSliderArea.x + levels.black_point * inputSliderArea.width;
                            const whiteX = inputSliderArea.x + levels.white_point * inputSliderArea.width;
                            const relativeRangeWidth = whiteX - blackX;

                            if (relativeRangeWidth < 1)
                                return;

                            let relativePos = (mouseX - blackX) / relativeRangeWidth;
                            relativePos = Math.max(0.0, Math.min(1.0, relativePos));
                            let invertedPos = 1.0 - relativePos;

                            if (invertedPos >= 0.5) {
                                levels.mid_point = 1.0 + ((invertedPos - 0.5) / 0.5) * (9.99 - 1.0);
                            }
                            else {
                                levels.mid_point = 0.01 + (invertedPos / 0.5) * (1.0 - 0.01);
                            }

                            levels.mid_point = parseFloat(levels.mid_point.toFixed(2));
                        }
                        needsUpdate = true;
                    }
                    // Handle output slider dragging
                    else if (this.dragging === 'output_black' || this.dragging === 'output_white') {
                        const outputSliderArea = this.getOutputSliderArea();

                        if (!outputSliderArea)
                            return;

                        let newLevel = ((mouseX - outputSliderArea.x) / outputSliderArea.width) * 255;
                        newLevel = Math.max(0, Math.min(255, newLevel));

                        if (this.dragging === 'output_black') {
                            levels.output_black = Math.min(newLevel, levels.output_white - 1);
                        }
                        else {
                            levels.output_white = Math.max(newLevel, levels.output_black + 1);
                        }
                        needsUpdate = true;
                    }


                    if (needsUpdate) {
                        updateWidgetsFromState();
                        this.setDirtyCanvas(true, true);
                        clearTimeout(this.updateTimeout);
                        this.updateTimeout = setTimeout(() => this.sendLevelUpdateToPython(), 50);
                    }
                },

                onMouseUp: function (event) {
                    if (this.dragging) {
                        this.dragging = null;
                        this.sendLevelUpdateToPython();
                    }
                }
            });

            const methodsToBind = ['getImageArea', 'getBarArea', 'getInputSliderArea', 'getOutputSliderArea', 'updateHistogram', 'sendLevelUpdateToPython', 'onDrawForeground', 'onMouseDown', 'onMouseMove', 'onMouseUp', 'drawHistogram', 'drawHandle'];
            methodsToBind.forEach(m => { if (node[m]) node[m] = node[m].bind(node); });
        }
    },
});



// Event listeners -----------

api.addEventListener("eses.levels_preview_base", ({ detail }) => {
    const node = app.graph.getNodeById(detail.node_id);

    if (node) {
        node.isManuallyResized = false;

        // Check for both auto_levels_updated 
        // and auto_color_updated
        if (detail.updated_levels) {
            
            // Runs for auto-levels or auto-color.
            // Update all widgets.
            node.allLevels = detail.updated_levels;

            const channelWidget = node.widgets.find(w => w.name === "channel");
            const channelMap = { "RGB": "rgb", "Red": "r", "Green": "g", "Blue": "b", "Mask": "mask" };

            node.activeChannel = channelMap[channelWidget.value] || "rgb";

            // Refresh the widgets for the currently
            // selected channel to show the newly reset
            // output levels.
            const blackWidget = node.widgets.find(w => w.name === "Black");
            const midWidget = node.widgets.find(w => w.name === "Mid");
            const whiteWidget = node.widgets.find(w => w.name === "White");
            const outputBlackWidget = node.widgets.find(w => w.name === "Output Black");
            const outputWhiteWidget = node.widgets.find(w => w.name === "Output White");

            const levels = node.allLevels[node.activeChannel];

            if (levels) {
                if (blackWidget)
                    blackWidget.value = parseFloat(levels.black_point.toFixed(2));

                if (midWidget)
                    midWidget.value = parseFloat(levels.mid_point.toFixed(2));

                if (whiteWidget)
                    whiteWidget.value = parseFloat(levels.white_point.toFixed(2));

                if (outputBlackWidget)
                    outputBlackWidget.value = Math.round(levels.output_black);

                if (outputWhiteWidget)
                    outputWhiteWidget.value = Math.round(levels.output_white);
            }
        }

        const onAssetLoad = async () => {
            if (!node.isManuallyResized) {
                const sizingObj = node.originalImage || node.originalMask;

                if (sizingObj) {
                    const aspectRatio = sizingObj.naturalWidth / sizingObj.naturalHeight;
                    const imageDrawWidth = node.size[0] - PADDING * 2;
                    const imageDrawHeight = imageDrawWidth / aspectRatio;
                    const interactiveHeight = HISTOGRAM_AREA_HEIGHT + OUTPUT_SLIDER_AREA_HEIGHT + AREA_SPACING;

                    node.size[1] = imageDrawHeight + interactiveHeight + HEADER_HEIGHT + PADDING * 2;
                    node.isManuallyResized = true; // Set to true after auto-resize
                }
            }
            await node.updateHistogram();
            await node.sendLevelUpdateToPython();
        };

        let assetsToLoad = 0;

        if (detail.image_data)
            assetsToLoad++;

        if (detail.mask_data)
            assetsToLoad++;

        if (assetsToLoad === 0) {
            onAssetLoad(); return;
        }

        let loadedCount = 0;
        const assetLoadedCallback = async () => {
            if (++loadedCount === assetsToLoad) {
                await onAssetLoad();
            }
        };

        if (detail.image_data) {
            const img = new Image();
            img.src = "data:image/png;base64," + detail.image_data;
            img.onload = assetLoadedCallback;
            node.originalImage = img;
        }
        else {
            node.originalImage = null;
        }

        if (detail.mask_data) {
            const mask = new Image();
            mask.src = "data:image/png;base64," + detail.mask_data;
            mask.onload = assetLoadedCallback;
            node.originalMask = mask;
        }
        else {
            node.originalMask = null;
        }
    }
});